﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace AdvancedCollections.TreeHelper
{
    public class AVLTreeNode<T> : BinaryTreeNode<T>
    {
        private int _depth;

        public AVLTreeNode(AVLTreeNode<T> parent, IComparer<T> comparer)
            : base(parent, comparer)
        {
        }

        public new AVLTreeNode<T> LeftChild
        {
            get { return (AVLTreeNode<T>)base.LeftChild; }
            protected set { base.LeftChild = value; }
        }

        public new AVLTreeNode<T> Parent
        {
            get { return (AVLTreeNode<T>)base.Parent; }
            protected set { base.Parent = value; }
        }

        public new AVLTreeNode<T> RightChild
        {
            get { return (AVLTreeNode<T>)base.RightChild; }
            protected set { base.RightChild = value; }
        }

        private int Balance { get { return LeftChild._depth - RightChild._depth; } }

        public new bool Add(T item, ref AVLTreeNode<T> root)
        {
            var node = FindNodeOrPlace(this, item);

            if (!node.IsEmpty)
                return false;
            else
            {
                node.FillNode(item);
                root = node.CheckBalance(root);
                return true;
            }
        }

        public new bool Remove(T item, ref AVLTreeNode<T> root)
        {
            var node = FindNodeOrPlace(this, item);

            if (node.IsEmpty)
                return false;
            else
            {
                if (node == root)
                {
                    root = (AVLTreeNode<T>)node.RemoveNodeAndReturnSuccessor();
                    root.Parent = null;
                }
                else
                    node.RemoveNodeAndReturnSuccessor();

                root = node.CheckBalance(root);
                return true;
            }
        }

        protected override void FillNode(T data)
        {
            base.FillNode(data);
            RecalculateDepth();
        }

        protected new AVLTreeNode<T> FindNodeOrPlace(AVLTreeNode<T> current, T item)
        {
            return (AVLTreeNode<T>)base.FindNodeOrPlace(current, item);
        }

        protected override void InitializeEmptyChildren()
        {
            LeftChild = new AVLTreeNode<T>(this, Comparer);
            RightChild = new AVLTreeNode<T>(this, Comparer);
        }

        private AVLTreeNode<T> CheckBalance(AVLTreeNode<T> root)
        {
            var currentNode = this;

            while (currentNode != null)
            {
                if (currentNode.Balance < -1 || currentNode.Balance > 1)
                    return currentNode.ReBalance(root);
                else
                    currentNode = currentNode.Parent;
            }

            return root;
        }

        private int GetMaxDepthFromChildren()
        {
            if (LeftChild._depth >= RightChild._depth)
                return LeftChild._depth;
            else
                return RightChild._depth;
        }

        /// <summary>
        ///
        /// </summary>
        /// <param name="root"></param>
        /// <returns>The new Root if changed otherwise the old Root</returns>
        private AVLTreeNode<T> LeftRotation(AVLTreeNode<T> root)
        {
            var newLeftChild = LeftChild.RightChild;

            if (Parent != null)
                Parent.SwitchSuccessor(this, LeftChild);
            else
            {
                root = LeftChild;
                LeftChild.Parent = null;
            }

            LeftChild.RightChild = this;
            Parent = LeftChild;
            LeftChild = newLeftChild;
            LeftChild.Parent = this;
            return root;
        }

        /// <summary>
        ///
        /// </summary>
        /// <param name="root"></param>
        /// <returns>The new Root if changed otherwise the old Root</returns>
        private AVLTreeNode<T> ReBalance(AVLTreeNode<T> root)
        {
            if (Balance < -1)
            {
                root = RightRotation(root);
            }
            else if (Balance > 1)
            {
                root = LeftRotation(root);
            }

            RecalculateDepth();

            return root;
        }

        private void RecalculateDepth()
        {
            var currentNode = this;
            int oldDepth;

            while (currentNode != null)
            {
                oldDepth = currentNode._depth;
                currentNode._depth = currentNode.GetMaxDepthFromChildren() + 1;

                if (oldDepth == currentNode._depth)
                    break;
                else
                    currentNode = currentNode.Parent;
            }
        }

        /// <summary>
        ///
        /// </summary>
        /// <param name="root"></param>
        /// <returns>The new Root if changed otherwise the old Root</returns>
        private AVLTreeNode<T> RightRotation(AVLTreeNode<T> root)
        {
            var newRightChild = RightChild.LeftChild;

            if (Parent != null)
                Parent.SwitchSuccessor(this, RightChild);
            else
            {
                root = RightChild;
                RightChild.Parent = null;
            }

            RightChild.LeftChild = this;
            Parent = RightChild;
            RightChild = newRightChild;
            RightChild.Parent = this;
            return root;
        }
    }
}